#!/bin/bash
# liho-sandbox: root-only helper the host drivers call inside the sandbox.
#
#   liho-sandbox mount <source> <guest-dir> ro|rw   mount a host folder under /mnt/liho
#   liho-sandbox mount-work <source>                mount the host work root (chat work folders) at /work
#   liho-sandbox umount-all                         unmount /work and everything under /mnt/liho
#   liho-sandbox run <id> <outer-port> <inner-port> <gw> <proxy-port> <relay-port> <bridge-port> -- <cmd> [args...]
#                                                   run <cmd> as `agent` in private net/mount/pid namespaces
#   liho-sandbox stop <id>                          stop a `run` and everything it started
#   liho-sandbox stop-all                           stop every `run` (those left by an earlier app instance too)
#   liho-sandbox egress <gateway-ip> <port>...      iptables: outbound only to these host ports (dedicated VMs only)
#   liho-sandbox egress-off                         restore the default ACCEPT policy
#   liho-sandbox version                            print the helper version
#
# Why `run` and not iptables on WSL: every WSL2 distribution shares ONE network
# namespace, so iptables rules here would also cut off Docker Desktop and the
# user's own distros, and owner matches on uid 1000 would hit their users.
# `run` gives the agent its own network namespace with only a loopback
# interface; the host proxy, the router relay and the host bridge are reachable
# through unix-socket bridges on fixed loopback ports, and opencode's port is
# bridged back out to 127.0.0.1:<outer-port> (WSL forwards it to the host).
#
# No eval, every variable quoted, every argument validated.
set -euo pipefail

HELPER_VERSION=8
LIHO_ROOT="/mnt/liho"
WORK_ROOT="/work"
RUN_ROOT="/run/liho"
# opencode's global config home (XDG_CONFIG_HOME): root-owned, so nothing the agent writes there is read back.
OPENCODE_CONFIG_HOME="/opt/liho/config"
AGENT_UID=1000
AGENT_GID=1000
PROXY_PORT_IN=3128
RELAY_PORT_IN=3129
BRIDGE_PORT_IN=3130
AGENT_PATH="/opt/liho/bin:/opt/liho/venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"

die() { printf 'liho-sandbox: %s\n' "$*" >&2; exit 2; }
usage() { sed -n '2,14p' "$0" >&2; exit 2; }

[[ "${EUID}" -eq 0 ]] || die "must run as root"

valid_guest_dir() {
    local dir="$1"
    [[ "${dir}" =~ ^/mnt/liho/[A-Za-z0-9._-]+$ ]] || return 1
    [[ "${dir}" != *".."* ]] || return 1
    return 0
}

valid_ipv4() {
    local ip="$1" octet
    [[ "${ip}" =~ ^([0-9]{1,3})\.([0-9]{1,3})\.([0-9]{1,3})\.([0-9]{1,3})$ ]] || return 1
    for octet in "${BASH_REMATCH[@]:1}"; do
        (( octet <= 255 )) || return 1
    done
    return 0
}

valid_port() {
    local port="$1"
    [[ "${port}" =~ ^[0-9]{1,5}$ ]] || return 1
    (( port >= 1 && port <= 65535 )) || return 1
    return 0
}

# 0 means "not used".
valid_port_or_zero() {
    [[ "$1" == "0" ]] && return 0
    valid_port "$1"
}

valid_id() {
    [[ "$1" =~ ^[A-Za-z0-9_-]{4,64}$ ]]
}

# A drvfs mount follows a symlink at its target, so a mount point swapped for a link
# between sessions would mount a host folder over any directory — /opt/liho/bin, whose
# scripts root runs, included. The root owns every component of a mount point for that
# reason: the point, its parent and the mount root must be real directories there, and
# the resolved path must be the path itself.
assert_mount_point() {
    local path="$1" resolved
    [[ ! -L "${path}" ]] || die "refusing a symlinked mount point: ${path}"
    [[ -d "${path}" ]] || die "mount point is not a directory: ${path}"
    resolved="$(realpath -e -- "${path}")" || die "cannot resolve the mount point: ${path}"
    [[ "${resolved}" == "${path}" ]] || die "mount point resolves elsewhere: ${path} -> ${resolved}"
}

cmd_mount() {
    [[ $# -eq 3 ]] || usage
    local src="$1" dst="$2" mode="$3" opts
    [[ -n "${src}" && "${src}" != *$'\n'* ]] || die "empty or invalid source"
    valid_guest_dir "${dst}" || die "guest dir must be ${LIHO_ROOT}/<name>: ${dst}"
    case "${mode}" in
        ro) opts="metadata,uid=${AGENT_UID},gid=${AGENT_GID},umask=022,ro" ;;
        rw) opts="metadata,uid=${AGENT_UID},gid=${AGENT_GID},umask=022" ;;
        *)  die "mode must be ro or rw" ;;
    esac
    mkdir -p -- "${dst}"
    assert_mount_point "${LIHO_ROOT}"
    assert_mount_point "${dst}"
    if mountpoint -q -- "${dst}"; then
        printf 'already mounted: %s\n' "${dst}"
        return 0
    fi
    # drvfs exists only on WSL; on Lima/Podman the folder arrives pre-mounted (handled above).
    mount -t drvfs -o "${opts}" -- "${src}" "${dst}"
    printf 'mounted %s -> %s (%s)\n' "${src}" "${dst}" "${mode}"
}

# The chats' work folders live on the host (so the user can open what the agent made);
# the whole root is one read-write drvfs mount, each chat works in /work/<chat-id>.
cmd_mount_work() {
    [[ $# -eq 1 ]] || usage
    local src="$1"
    [[ -n "${src}" && "${src}" != *$'\n'* ]] || die "empty or invalid source"
    mkdir -p -- "${WORK_ROOT}"
    assert_mount_point "${WORK_ROOT}"
    if mountpoint -q -- "${WORK_ROOT}"; then umount -l -- "${WORK_ROOT}"; fi
    mount -t drvfs -o "metadata,uid=${AGENT_UID},gid=${AGENT_GID},umask=022" -- "${src}" "${WORK_ROOT}"
    printf 'mounted %s -> %s (rw)\n' "${src}" "${WORK_ROOT}"
}

cmd_umount_all() {
    [[ $# -eq 0 ]] || usage
    local target
    if mountpoint -q -- "${WORK_ROOT}"; then umount -l -- "${WORK_ROOT}" && printf 'unmounted %s\n' "${WORK_ROOT}"; fi
    while IFS= read -r target; do
        [[ -n "${target}" ]] || continue
        umount -l -- "${target}" && printf 'unmounted %s\n' "${target}"
    done < <(findmnt -rn -o TARGET | grep -E "^${LIHO_ROOT}/" | awk '{ print length($0) " " $0 }' | sort -rn | cut -d' ' -f2-)
    find "${LIHO_ROOT}" -mindepth 1 -maxdepth 1 -type d -empty -exec rmdir -- {} + 2>/dev/null || true
}

# ---- run: outer half (root namespace of the distribution) -------------------

cmd_run() {
    [[ $# -ge 9 ]] || usage
    local id="$1" outer="$2" inner="$3" gw="$4" proxy="$5" relay="$6" bridge="$7"
    shift 7
    [[ "$1" == "--" ]] || usage
    shift
    [[ $# -ge 1 ]] || usage
    valid_id "${id}" || die "invalid run id"
    valid_port_or_zero "${outer}" || die "invalid outer port"
    valid_port_or_zero "${inner}" || die "invalid inner port"
    valid_port_or_zero "${proxy}" || die "invalid proxy port"
    valid_port_or_zero "${relay}" || die "invalid relay port"
    valid_port_or_zero "${bridge}" || die "invalid bridge port"
    if [[ "${proxy}" != 0 || "${relay}" != 0 || "${bridge}" != 0 ]]; then
        valid_ipv4 "${gw}" || die "invalid gateway address"
    fi

    local dir="${RUN_ROOT}/${id}"
    mkdir -p -m 700 -- "${RUN_ROOT}"
    mkdir -p -m 700 -- "${dir}"
    chmod 700 -- "${dir}"
    printf '%s\n' "$$" > "${dir}/run.pid"
    # opencode creates its config dir at start: it exists, and only root writes to it.
    install -d -o root -g root -m 755 -- "${OPENCODE_CONFIG_HOME}" "${OPENCODE_CONFIG_HOME}/opencode"

    local -a pids=()
    [[ "${proxy}" != 0 ]] && { socat "UNIX-LISTEN:${dir}/proxy.sock,fork,mode=600" "TCP:${gw}:${proxy}" & pids+=("$!"); }
    [[ "${relay}" != 0 ]] && { socat "UNIX-LISTEN:${dir}/relay.sock,fork,mode=600" "TCP:${gw}:${relay}" & pids+=("$!"); }
    [[ "${bridge}" != 0 ]] && { socat "UNIX-LISTEN:${dir}/bridge.sock,fork,mode=600" "TCP:${gw}:${bridge}" & pids+=("$!"); }
    [[ "${outer}" != 0 ]] && { socat "TCP-LISTEN:${outer},bind=127.0.0.1,fork,reuseaddr" "UNIX-CONNECT:${dir}/oc.sock" & pids+=("$!"); }

    # Background jobs of a non-interactive shell get /dev/null as stdin; `<&0` keeps the
    # caller's stdin for the agent command (the app streams attachments through it).
    local child=""
    cleanup() {
        [[ -n "${child}" ]] && kill -TERM "${child}" 2>/dev/null || true
        ((${#pids[@]})) && kill "${pids[@]}" 2>/dev/null || true
        rm -rf -- "${dir}"
    }
    trap cleanup EXIT
    trap 'kill -TERM "${child}" 2>/dev/null || true' TERM INT HUP

    # A socat bridge that dies at once (its port is taken) fails the run here: the host would
    # otherwise face an unexplained ECONNREFUSED on the port it was promised.
    if ((${#pids[@]})); then
        sleep 0.3
        local p=''
        for p in "${pids[@]}"; do
            if ! kill -0 "${p}" 2>/dev/null; then
                echo "liho-sandbox: a bridge listener exited at start (is port ${outer} taken?)" >&2
                exit 1
            fi
        done
    fi

    unshare --net --mount --pid --fork --kill-child --mount-proc --propagation private -- \
        "$0" inner "${id}" "${inner}" "${proxy}" "${relay}" "${bridge}" -- "$@" <&0 &
    child="$!"
    local code=0
    wait "${child}" || code=$?
    exit "${code}"
}

# ---- run: inner half (PID 1 of the private namespaces, still root) -----------

cmd_inner() {
    [[ $# -ge 7 ]] || usage
    local id="$1" inner="$2" proxy="$3" relay="$4" bridge="$5"
    shift 5
    [[ "$1" == "--" ]] || usage
    shift
    valid_id "${id}" || die "invalid run id"
    local dir="${RUN_ROOT}/${id}"

    ip link set lo up

    # A private /tmp and nothing that reaches the Windows desktop or other distros:
    # /mnt/wsl is shared between all WSL distributions, /mnt/wslg and the X11/Wayland
    # sockets reach the Windows desktop, /run/WSL holds the interop sockets.
    mount -t tmpfs -o mode=1777,size=2g tmpfs /tmp
    local hidden
    for hidden in /mnt/wsl /mnt/wslg /run/WSL /run/user; do
        [[ -d "${hidden}" ]] && mount -t tmpfs -o mode=755,size=64k tmpfs "${hidden}" 2>/dev/null || true
    done

    # The host wrote the runtime environment (NUL-separated) into a root-only file; read and delete it.
    local -a envv=()
    if [[ -f "${dir}/env" ]]; then
        mapfile -d '' -t envv < "${dir}/env"
        rm -f -- "${dir}/env"
    fi

    [[ "${proxy}" != 0 ]] && socat "TCP-LISTEN:${PROXY_PORT_IN},bind=127.0.0.1,fork,reuseaddr" "UNIX-CONNECT:${dir}/proxy.sock" &
    [[ "${relay}" != 0 ]] && socat "TCP-LISTEN:${RELAY_PORT_IN},bind=127.0.0.1,fork,reuseaddr" "UNIX-CONNECT:${dir}/relay.sock" &
    [[ "${bridge}" != 0 ]] && socat "TCP-LISTEN:${BRIDGE_PORT_IN},bind=127.0.0.1,fork,reuseaddr" "UNIX-CONNECT:${dir}/bridge.sock" &
    [[ "${inner}" != 0 ]] && socat "UNIX-LISTEN:${dir}/oc.sock,fork,mode=600" "TCP:127.0.0.1:${inner}" &

    cd /work 2>/dev/null || cd /

    setpriv --reuid="${AGENT_UID}" --regid="${AGENT_GID}" --init-groups --no-new-privs \
        --inh-caps=-all --bounding-set=-all -- \
        /usr/bin/env -i \
        HOME=/home/agent USER=agent LOGNAME=agent SHELL=/bin/bash \
        PATH="${AGENT_PATH}" LANG=C.UTF-8 LC_ALL=C.UTF-8 TMPDIR=/tmp \
        PLAYWRIGHT_BROWSERS_PATH=/opt/liho/ms-playwright \
        NODE_PATH=/usr/lib/node_modules:/usr/lib/node_modules/@playwright/mcp/node_modules \
        OPENCODE_DISABLE_AUTOUPDATE=1 OPENCODE_DISABLE_LSP_DOWNLOAD=1 \
        "${envv[@]}" "$@" <&0 &
    local agent="$!"
    trap 'kill -TERM "${agent}" 2>/dev/null || true' TERM INT HUP
    local code=0
    wait "${agent}" || code=$?
    # PID 1 exiting tears the namespace down, socats included.
    exit "${code}"
}

cmd_stop() {
    [[ $# -eq 1 ]] || usage
    valid_id "$1" || die "invalid run id"
    local dir="${RUN_ROOT}/$1" pid
    [[ -f "${dir}/run.pid" ]] || { printf 'not running: %s\n' "$1"; return 0; }
    pid="$(cat -- "${dir}/run.pid")"
    [[ "${pid}" =~ ^[0-9]+$ ]] || die "bad pid file"
    # PID files outlive a SIGKILL and the kernel may reuse their number. Pin the
    # process before checking its command line, then signal only through that fd.
    # This also handles runs started by older helpers without a new PID format.
    python3 -I - "${pid}" "$0" "$1" <<'PY'
import os
from pathlib import Path
import select
import signal
import sys

pid = int(sys.argv[1])
if pid <= 1:
    raise SystemExit('liho-sandbox: bad pid file')
try:
    descriptor = os.pidfd_open(pid)
except ProcessLookupError:
    raise SystemExit(0)
try:
    try:
        command = Path(f'/proc/{pid}/cmdline').read_bytes().split(b'\0')
        expected = [os.fsencode(sys.argv[2]), b'run', os.fsencode(sys.argv[3])]
        if len(command) < 4 or Path(os.fsdecode(command[0])).name != 'bash' or command[1:4] != expected:
            raise SystemExit(0)
        signal.pidfd_send_signal(descriptor, signal.SIGTERM)
        ready, _, _ = select.select([descriptor], [], [], 3)
        if not ready:
            signal.pidfd_send_signal(descriptor, signal.SIGKILL)
    except (ProcessLookupError, FileNotFoundError):
        pass
finally:
    os.close(descriptor)
PY
    rm -rf -- "${dir}"
    printf 'stopped %s\n' "$1"
}

# Runs of an app instance that crashed keep their namespaces (and mounts as they were): the
# first start of a new instance stops them all.
cmd_stop_all() {
    [[ $# -eq 0 ]] || usage
    local dir id
    [[ -d "${RUN_ROOT}" ]] || return 0
    for dir in "${RUN_ROOT}"/*/; do
        [[ -d "${dir}" ]] || continue
        id="$(basename -- "${dir}")"
        valid_id "${id}" || continue
        cmd_stop "${id}"
    done
}

ipt() { iptables -w 5 "$@"; }
# IPv6 fails closed: without working ip6tables the interfaces lose IPv6 instead.
ipt6_closed() {
    if command -v ip6tables >/dev/null 2>&1 && ip6tables -w 5 -F OUTPUT && ip6tables -w 5 -P OUTPUT DROP && ip6tables -w 5 -A OUTPUT -o lo -j ACCEPT; then
        return 0
    fi
    [[ -e /proc/sys/net/ipv6/conf/all/disable_ipv6 ]] || return 0
    sysctl -q -w net.ipv6.conf.all.disable_ipv6=1 net.ipv6.conf.default.disable_ipv6=1 >/dev/null || die "IPv6 egress cannot be restricted"
}
ipt6_open() {
    if command -v ip6tables >/dev/null 2>&1; then ip6tables -w 5 -P OUTPUT ACCEPT 2>/dev/null || true; ip6tables -w 5 -F OUTPUT 2>/dev/null || true; fi
}

# Outbound traffic only to the given host ports. No DNS: the host proxy resolves names, and a
# resolver on the gateway would be a channel out of every network mode.
cmd_egress() {
    [[ $# -ge 2 ]] || usage
    local gw="$1" port
    shift
    valid_ipv4 "${gw}" || die "invalid gateway IPv4 address: ${gw}"
    for port in "$@"; do valid_port "${port}" || die "invalid port: ${port}"; done
    ipt -F OUTPUT
    ipt -P OUTPUT DROP
    ipt -A OUTPUT -o lo -j ACCEPT
    ipt -A OUTPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    # The proxy (unless the network is off), the router relay and (with computer use) the host bridge.
    for port in "$@"; do ipt -A OUTPUT -p tcp -d "${gw}" --dport "${port}" -j ACCEPT; done
    ipt6_closed
    printf 'egress restricted to %s ports %s\n' "${gw}" "$*"
}

cmd_egress_off() {
    [[ $# -eq 0 ]] || usage
    ipt -P OUTPUT ACCEPT
    ipt -F OUTPUT
    ipt6_open
    printf 'egress unrestricted\n'
}

[[ $# -ge 1 ]] || usage
command="$1"
shift
case "${command}" in
    mount)      cmd_mount "$@" ;;
    mount-work) cmd_mount_work "$@" ;;
    umount-all) cmd_umount_all "$@" ;;
    run)        cmd_run "$@" ;;
    inner)      cmd_inner "$@" ;;
    stop)       cmd_stop "$@" ;;
    stop-all)   cmd_stop_all "$@" ;;
    egress)     cmd_egress "$@" ;;
    egress-off) cmd_egress_off "$@" ;;
    version)    printf '%s\n' "${HELPER_VERSION}" ;;
    *)          usage ;;
esac
